package com.test.agent.service;

import com.alibaba.fastjson.JSON;
import com.mongodb.client.gridfs.GridFSBucket;
import com.mongodb.client.gridfs.model.GridFSFile;
import com.test.agent.Config;
import com.test.common.entity.perf.PerfAgentEntity;
import com.test.common.entity.perf.PerfReportEntity;
import com.test.common.entity.perf.PerfResourceEntity;
import com.test.common.entity.perf.PerfScriptEntity;
import com.test.common.repository.perf.PerfReportRepository;
import com.test.common.repository.perf.PerfScriptRepository;
import com.test.common.response.BizException;
import com.test.common.utils.CommandUtils;
import com.test.common.utils.ScriptUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.io.FileUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.gridfs.GridFsCriteria;
import org.springframework.data.mongodb.gridfs.GridFsResource;
import org.springframework.data.mongodb.gridfs.GridFsTemplate;
import org.springframework.http.*;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;

import javax.annotation.Resource;
import java.io.File;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.HashMap;
import java.util.List;
import java.util.Optional;

@Slf4j
@Service
public class AgentService {

    @Resource
    private Config config;

    @Resource
    private CommandUtils commandUtils;

    @Resource
    private ScriptUtils scriptUtils;

    @Resource
    private PerfScriptRepository perfScriptRepository;

    @Resource
    private PerfReportRepository perfReportRepository;

    @Resource
    private GridFsTemplate gridFsTemplate;

    @Resource
    private GridFSBucket gridFSBucket;

    private final RestTemplate restTemplate = new RestTemplate();

    public HashMap<String, Object> reflushStatus() {

        int threads = config.getMaxThreads();

        // 查询当前正在执行的脚本
        String[] scriptNames = commandUtils
                .exec("ps -ef | grep ApacheJMeter.jar | grep -v grep | grep -v 'ps -ef' | awk '{print $NF}'")
                .split("\n");

        // TODO 获取这个脚本的并发设置

        return new HashMap<String, Object>() {{
            put("threads", threads);
            put("processes",
                    commandUtils.exec("ps -ef | grep ApacheJMeter.jar | grep -v grep | grep -v 'ps -ef'").split("\n"));
        }};
    }

    public String startTask(int reportId) {

        // 只能在linux下执行
        if (!System.getProperty("os.name").toLowerCase().contains("linux")) {
            throw new BizException("执行机暂时仅支持linux");
        }

        Optional<PerfReportEntity> reportEntityOptional = perfReportRepository.findById(reportId);
        if (reportEntityOptional.isPresent()) {

            PerfReportEntity perfReportEntity = reportEntityOptional.get();

            // 检查脚本是否正在执行
            String current = commandUtils.exec("ps -ef | grep " + perfReportEntity.getUuid() + " | grep -v grep");
            if (StringUtils.isNotEmpty(current)) {
                log.warn("警告：当前脚本已有其他执行进程：\n{}", current);
            }

            Optional<PerfScriptEntity> optionalPerfScriptEntity = perfScriptRepository.findById(perfReportEntity.getScriptId());
            if (optionalPerfScriptEntity.isPresent()) {

                PerfScriptEntity perfScriptEntity = optionalPerfScriptEntity.get();

                String scriptName;
                File scriptFile;
                File resourceFolder;
                switch (perfScriptEntity.getPlatform()) {
                    case JMETER:
                        // JMeter执行机，将脚本更名为报告中的uuid，方便解析器收集数据
                        scriptName = perfReportEntity.getUuid() + ".jmx";
                        scriptFile = new File(config.getJmeterHome() + File.separator + "scripts" + File.separator + scriptName);
                        resourceFolder = new File(config.getJmeterHome() + File.separator + "resources");
                        if (!resourceFolder.exists()) {
                            log.info("初始化资源文件夹：{}", resourceFolder.mkdirs());
                        }
                        break;
                    case GATLING:
                        // Gatling执行机，脚本不需要换名字
                        scriptName = perfScriptEntity.getName();
                        scriptFile = new File(config.getGatlingHome() + File.separator + "user-files" + File.separator + "simulations" + File.separator + scriptName);
                        resourceFolder = new File(config.getGatlingHome() + File.separator + "user-files" + File.separator + "resources");
                        if (!resourceFolder.exists()) {
                            log.info("初始化资源文件夹：{}", resourceFolder.mkdirs());
                        }
                        break;
                    default:
                        throw new BizException("不支持的脚本平台：" + perfScriptEntity.getPlatform());
                }

                // 下载脚本
                downScript(perfScriptEntity, scriptFile);

                // 下载资源文件
                downResource(perfScriptEntity, resourceFolder);

                // 修改脚本：修改启动时间、持续时间、并发等数据
                scriptUtils.modifyScript(perfScriptEntity.getPlatform(), scriptFile, perfReportEntity);

                // 执行脚本
                switch (perfScriptEntity.getPlatform()) {
                    case JMETER:
                        // Jmeter脚本执行只需要指定文件即可，依赖的资源文件写在脚本中
                        commandUtils.exec("nohup sh " + config.getJmeterHome() + "/bin/NewDriver -n -t " + scriptFile.getAbsolutePath() + " >> jmeter.echo 2>&1 &");
                        break;
                    case GATLING:
                        log.warn("Gatling需要手动配置好脚本目录和资源文件目录");
                        // 再执行
                        commandUtils.exec("nohup sh " + config.getGatlingHome() + "/bin/gatling.sh -s " + scriptName.replace(".scala", "") + " -rd " + perfReportEntity.getUuid() + " >> gatling.echo 2>&1 &");
                        break;
                    default:
                        throw new BizException("不支持的脚本平台：" + perfScriptEntity.getPlatform());
                }

                // 等待进程启动，Gatling会比较慢
                try {
                    Thread.sleep(5000);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }

                // 查询进程pid
                String pid = commandUtils.exec("ps -ef | grep " + scriptFile.getName() + " | grep -v grep | awk '{print $2}'").trim();

                if (pid.equals("")) {
                    throw new BizException("执行失败，没有查询到相关进程");
                } else {
                    return "执行成功，进程号为：" + pid;
                }
            } else {
                throw new BizException("脚本文件丢失：" + perfReportEntity.getScriptId());
            }
        } else {
            throw new BizException("找不到任务报告对象：" + reportId);
        }
    }

    private void downResource(PerfScriptEntity perfScriptEntity, File resourceFolder) {

        List<PerfResourceEntity> resources = perfScriptEntity.getResources();
        resources.forEach(perfResourceEntity -> {
            try {
                File resourceFile = new File(resourceFolder.getAbsolutePath() + File.separator + perfResourceEntity.getName());
                log.info("资源文件保存到：{}", resourceFile.getAbsolutePath());

                GridFSFile fsFile = gridFsTemplate.findOne(Query.query(GridFsCriteria.where("_id").is(perfResourceEntity.getObjectId())));
                if (fsFile != null) {
                    GridFsResource fsResource = new GridFsResource(fsFile, gridFSBucket.openDownloadStream(fsFile.getObjectId()));
                    log.info("文件大小|{}", fsResource.contentLength());
                    // 直接保存到资源文件目录
                    FileUtils.copyInputStreamToFile(fsResource.getInputStream(), resourceFile);
                } else {
                    throw new BizException("资源文件丢失：" + perfScriptEntity.getObjectId());
                }
            } catch (Exception e) {
                log.error("获取文件信息异常|{}", e.getMessage());
                e.printStackTrace();
                throw new BizException("获取文件信息异常", e);
            }
        });
    }

    private void downScript(PerfScriptEntity perfScriptEntity, File destFile) {

        try {
            GridFSFile fsFile = gridFsTemplate.findOne(Query.query(GridFsCriteria.where("_id").is(perfScriptEntity.getObjectId())));
            if (fsFile != null) {
                GridFsResource fsResource = new GridFsResource(fsFile, gridFSBucket.openDownloadStream(fsFile.getObjectId()));
                log.info("文件大小|{}", fsResource.contentLength());

                FileUtils.copyInputStreamToFile(fsResource.getInputStream(), destFile);
            } else {
                throw new BizException("资源文件丢失：" + perfScriptEntity.getObjectId());
            }
        } catch (Exception e) {
            log.error("获取文件信息异常|{}", e.getMessage());
            e.printStackTrace();
            throw new BizException("获取文件信息异常", e);
        }
    }

    public String stopTask(String reportUuid) {

        if (!System.getProperty("os.name").toLowerCase().contains("linux")) {
            throw new BizException("执行机暂时仅支持linux");
        }

        commandUtils.exec("ps -ef | grep " + reportUuid + " | grep -v grep | awk '{print $2}' | xargs kill -9");

        // 检查脚本是否正在执行
        if (StringUtils.isEmpty(commandUtils.exec("ps -ef | grep " + reportUuid + " | grep -v grep").trim())) {
            return "执行成功，进程已强制关闭";
        } else {
            throw new BizException("执行失败，请重试");
        }
    }

    // 获取本地地址
    private String currentUrl;

    {
        try {
            currentUrl = "http://" + InetAddress.getLocalHost().getHostAddress() + ":18090";
        } catch (UnknownHostException e) {
            e.printStackTrace();
        }
    }

    // 3s心跳一次，初次延时3s执行
    @Scheduled(fixedDelay = 10000, initialDelay = 3000)
    public void heartBeat() {

        HttpHeaders httpHeaders = new HttpHeaders();
        httpHeaders.add("content-type", "application/json");

        PerfAgentEntity perfAgentEntity = new PerfAgentEntity(currentUrl, new String[0]);

        HttpEntity<String> httpEntity = new HttpEntity<>(JSON.toJSONString(perfAgentEntity), httpHeaders);

        ResponseEntity<String> exchange = restTemplate.exchange(config.getWebServer(), HttpMethod.POST, httpEntity, String.class);

        if (exchange.getStatusCode().equals(HttpStatus.OK)) {
            log.info(exchange.getBody());
        } else {
            log.error(exchange.getBody());
        }
    }
}
